"""
Core client functionality for Dolphin MCP.
"""

import os
import sys
import json
import asyncio
import logging
from typing import Any, Dict, List, Optional, Union, AsyncGenerator

from contextlib import AsyncExitStack
from mcp.client.sse import sse_client
from mcp import ClientSession

from core.providers import provider

logger = logging.getLogger("weather_mcp_client")


class SSEMCPClient:
    """Implementation for a SSE-based MCP server."""

    def __init__(self, server_name: str, url: str, models_cfg, args=None, env=None):
        self.server_name = server_name
        self.args = args or []
        self.url = url
        self.tools = []

        self.exit_stack = AsyncExitStack()
        self.models_cfg = models_cfg
        
        self._streams_context = None
        self._session_context = None
        self.session = None
    async def start(self):
        """
        Connect to an SSE MCP server
        
        Args:
            server_script_path: Path to the server script (.py or .js)
        """
        try:
            sse_transport = await self.exit_stack.enter_async_context(sse_client(url=self.url))
            self.sse, self.write = sse_transport
            self.session = await self.exit_stack.enter_async_context(ClientSession(self.sse, self.write))
            # Initialize
            await self.session.initialize()
            # List available tools
            response = await self.session.list_tools()
            tools = response.tools
            self.tools=tools
            logger.info(f"Server {self.server_name}: Connected to SSE server with with tools: {[tool.name for tool in tools]}.")
            return True
        except Exception as e:
            logger.error(f"Server {self.server_name}: SSE connection error: {str(e)}")
            return False
    async def create_model(self, provider_name: str):
        self.model=provider.Provider(self.models_cfg).create(provider_name)
    def create_messages(self, query: str):
        # 4) Build conversation
        # Get system message - either from systemMessageFile, systemMessage, or default
        conversation = []
        chosen_model=self.models_cfg
        system_msg = "You are a weather forecast assistant with ability to use tool function call according to location or geographical coordinates."
        if "systemMessageFile" in chosen_model:
            try:
                with open(chosen_model["systemMessageFile"], "r", encoding="utf-8") as f:
                    system_msg = f.read()
            except Exception as e:
                logger.warning(f"Failed to read system message file: {e}")
                # Fall back to direct systemMessage if available
                conversation.append({"role": "system", "content": chosen_model.get("systemMessage", system_msg)})
        else:
            conversation.append({"role": "system", "content": chosen_model.get("systemMessage", system_msg)})
        if "systemMessageFiles" in chosen_model:
            for file in chosen_model["systemMessageFiles"]:
                try:
                    with open(file, "r", encoding="utf-8") as f:
                        system_msg = f.read()
                        conversation.append({"role": "system", "content": "File: " + file + "\n" + system_msg})
                except Exception as e:
                    logger.warning(f"Failed to read system message file: {e}")
 
        conversation.append({"role": "user", "content": query})
        return conversation
    async def process_query(self, query: str) -> str:
        """Process a query using Claude and available tools"""
        messages = [
            {
                "role": "user",
                "content": query
            }
        ]

        # gather tools with type mcp.Tool
        all_functions = []
        logger.info(f"Server {self.server_name}: Gather tools in process query...")
        for t in self.tools:
            input_schema = t.inputSchema or {"type": "object", "properties": {}}
            fn_def = {
                "name": t.name,
                "description": t.description,
                "parameters": input_schema
            }
            all_functions.append(fn_def)
       
        #logger.info(f"Server {self.server_name}: list all_function info {json.dumps(all_functions, indent=2)}")

        # Initial Claude API call
        """
        response = self.anthropic.messages.create(
            #model="claude-3-5-sonnet-20241022",
            model="qwen3:0.6b",
            max_tokens=1000,
            messages=messages,
            tools=available_tools
        )
        """
        conversation=self.create_messages(query)
        response=await self.model.generate_response(
            conversation,
            all_functions=all_functions
        )
        # Process response and handle tool calls
        final_text = []

        for content in response.content:
            logger.debug(f"response content is {content}")
            if content.type == 'text':
                final_text.append(content.text)
            elif content.type == 'tool_use':
                tool_name = content.name
                arguments = content.input
                
                # Execute tool call
                result = await self.session.call_tool(tool_name, arguments)
                logger.info(f"[Calling tool={tool_name} with arguments={arguments}] with result: [{result.content[0].text}]")
                final_text.append(f"[Calling tool {tool_name} with args {arguments}]")

                # Continue conversation with tool results
                if hasattr(content, 'text') and content.text:
                    messages.append({
                      "role": "assistant",
                      "content": content.text
                    })
                messages.append({
                    "role": "user", 
                    "content": result.content[0].text
                })
                # Get next response from Claude
                response = await self.model.generate_response(
                    messages, 
                    all_functions=None,
                    max_round=1,
                )

                final_text.append(response.content[0].text)

        return "\n".join(final_text)
        
    async def chat_loop(self):
        """Run an interactive chat loop"""
        print("\nMCP Client Started!")
        print("Type your queries or 'quit' to exit.")
        
        while True:
            try:
                query = input("\nQuery: ").strip()
                
                if query.lower() == 'quit':
                    break
                    
                response = await self.process_query(query)
                print("\n" + response)
                    
            except Exception as e:
                print(f"\nError: {str(e)}")
    
    async def cleanup(self):
        """Clean up resources"""
        await self.exit_stack.aclose()


class MCPClient:
    """Implementation for a single MCP server."""
    def __init__(self, server_name, command, args=None, env=None):
        self.server_name = server_name
        self.command = command
        self.args = args or []
        print(self.args)
        self.env = env
        self.process = None
        self.tools = []
        self.request_id = 0
        self.protocol_version = "2024-11-05"
        self.receive_task = None
        self.responses = {}
        self.server_capabilities = {}
        self._shutdown = False
        self._cleanup_lock = asyncio.Lock()

    async def _receive_loop(self):
        if not self.process or self.process.stdout.at_eof():
            return
        try:
            while not self.process.stdout.at_eof():
                line = await self.process.stdout.readline()
                if not line:
                    break
                try:
                    message = json.loads(line.decode().strip())
                    self._process_message(message)
                except json.JSONDecodeError:
                    pass
        except Exception:
            pass

    def _process_message(self, message: dict):
        if "jsonrpc" in message and "id" in message:
            if "result" in message or "error" in message:
                self.responses[message["id"]] = message
            else:
                # request from server, not implemented
                resp = {
                    "jsonrpc": "2.0",
                    "id": message["id"],
                    "error": {
                        "code": -32601,
                        "message": f"Method {message.get('method')} not implemented in client"
                    }
                }
                asyncio.create_task(self._send_message(resp))
        elif "jsonrpc" in message and "method" in message and "id" not in message:
            # notification from server
            pass

    async def start(self):
        expanded_args = []
        for a in self.args:
            if isinstance(a, str) and "~" in a:
                expanded_args.append(os.path.expanduser(a))
            else:
                expanded_args.append(a)

        env_vars = os.environ.copy()
        if self.env:
            env_vars.update(self.env)
        print(self.args)
        logger.info(f"Server {self.server_name}: expanded_args and env_vars are prepared")
        try:
            self.process = await asyncio.create_subprocess_exec(
                self.command,
                *expanded_args,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=env_vars
            )
            logger.info(f"Server {self.server_name}: create subprocess exec ")
            self.receive_task = asyncio.create_task(self._receive_loop())
            logger.info(f"Server {self.server_name}: create subprocess exec task")
            return await self._perform_initialize()
        except Exception:
            return False

    async def _perform_initialize(self):
        self.request_id += 1
        req_id = self.request_id
        req = {
            "jsonrpc": "2.0",
            "id": req_id,
            "method": "initialize",
            "params": {
                "protocolVersion": self.protocol_version,
                "capabilities": {"sampling": {}},
                "clientInfo": {
                    "name": "DolphinMCPClient",
                    "version": "1.0.0"
                }
            }
        }
        await self._send_message(req)

        start = asyncio.get_event_loop().time()
        timeout = 10  # Increased timeout to 10 seconds
        logger.info(f"Server {self.server_name}: start asyncio.get_event_loop")
        while asyncio.get_event_loop().time() - start < timeout:
            if req_id in self.responses:
                resp = self.responses[req_id]
                del self.responses[req_id]
                if "error" in resp:
                    logger.error(f"Server {self.server_name}: Initialize error: {resp['error']}")
                    return False
                if "result" in resp:
                    elapsed = asyncio.get_event_loop().time() - start
                    logger.info(f"Server {self.server_name}: Initialized in {elapsed:.2f}s")
                    note = {"jsonrpc": "2.0", "method": "notifications/initialized"}
                    await self._send_message(note)
                    init_result = resp["result"]
                    self.server_capabilities = init_result.get("capabilities", {})
                    return True
            await asyncio.sleep(0.05)
        logger.error(f"Server {self.server_name}: Initialize timed out after {timeout}s")
        return False

    async def list_tools(self):
        if not self.process:
            return []
        self.request_id += 1
        rid = self.request_id
        req = {
            "jsonrpc": "2.0",
            "id": rid,
            "method": "tools/list",
            "params": {}
        }
        await self._send_message(req)

        start = asyncio.get_event_loop().time()
        timeout = 10  # Increased timeout to 10 seconds
        while asyncio.get_event_loop().time() - start < timeout:
            if rid in self.responses:
                resp = self.responses[rid]
                del self.responses[rid]
                if "error" in resp:
                    logger.error(f"Server {self.server_name}: List tools error: {resp['error']}")
                    return []
                if "result" in resp and "tools" in resp["result"]:
                    elapsed = asyncio.get_event_loop().time() - start
                    logger.info(f"Server {self.server_name}: Listed {len(resp['result']['tools'])} tools in {elapsed:.2f}s")
                    self.tools = resp["result"]["tools"]
                    return self.tools
            await asyncio.sleep(0.05)
        logger.error(f"Server {self.server_name}: List tools timed out after {timeout}s")
        return []

    async def call_tool(self, tool_name: str, arguments: dict):
        if not self.process:
            return {"error": "Not started"}
        self.request_id += 1
        rid = self.request_id
        req = {
            "jsonrpc": "2.0",
            "id": rid,
            "method": "tools/call",
            "params": {
                "name": tool_name,
                "arguments": arguments
            }
        }
        await self._send_message(req)

        start = asyncio.get_event_loop().time()
        timeout = 3600  # Increased timeout to 30 seconds
        while asyncio.get_event_loop().time() - start < timeout:
            if rid in self.responses:
                resp = self.responses[rid]
                del self.responses[rid]
                if "error" in resp:
                    logger.error(f"Server {self.server_name}: Tool {tool_name} error: {resp['error']}")
                    return {"error": resp["error"]}
                if "result" in resp:
                    elapsed = asyncio.get_event_loop().time() - start
                    logger.info(f"Server {self.server_name}: Tool {tool_name} completed in {elapsed:.2f}s")
                    return resp["result"]
            await asyncio.sleep(0.01)  # Reduced sleep interval for more responsive streaming
            if asyncio.get_event_loop().time() - start > 5:  # Log warning after 5 seconds
                logger.warning(f"Server {self.server_name}: Tool {tool_name} taking longer than 5s...")
        logger.error(f"Server {self.server_name}: Tool {tool_name} timed out after {timeout}s")
        return {"error": f"Timeout waiting for tool result after {timeout}s"}

    async def _send_message(self, message: dict):
        if not self.process or self._shutdown:
            logger.error(f"Server {self.server_name}: Cannot send message - process not running or shutting down")
            return False
        try:
            data = json.dumps(message) + "\n"
            self.process.stdin.write(data.encode())
            await self.process.stdin.drain()
            return True
        except Exception as e:
            logger.error(f"Server {self.server_name}: Error sending message: {str(e)}")
            return False

    async def stop(self):
        async with self._cleanup_lock:
            if self._shutdown:
                return
            self._shutdown = True

            if self.receive_task and not self.receive_task.done():
                self.receive_task.cancel()
                try:
                    await self.receive_task
                except asyncio.CancelledError:
                    pass

            if self.process:
                try:
                    # Try to send a shutdown notification first
                    try:
                        note = {"jsonrpc": "2.0", "method": "shutdown"}
                        await self._send_message(note)
                        # Give a small window for the process to react
                        await asyncio.sleep(0.5)
                    except:
                        pass

                    # Close stdin before terminating to prevent pipe errors
                    if self.process.stdin:
                        self.process.stdin.close()

                    # Try graceful shutdown first
                    self.process.terminate()
                    try:
                        # Use a shorter timeout to make cleanup faster
                        await asyncio.wait_for(self.process.wait(), timeout=1.0)
                    except asyncio.TimeoutError:
                        # Force kill if graceful shutdown fails
                        logger.warning(f"Server {self.server_name}: Force killing process after timeout")
                        self.process.kill()
                        try:
                            await asyncio.wait_for(self.process.wait(), timeout=1.0)
                        except asyncio.TimeoutError:
                            logger.error(f"Server {self.server_name}: Process did not respond to SIGKILL")
                except Exception as e:
                    logger.error(f"Server {self.server_name}: Error during process cleanup: {str(e)}")
                finally:
                    # Make sure we clear the reference
                    self.process = None

    # Alias close to stop for backward compatibility
    async def close(self):
        await self.stop()

    # Add async context manager support
    async def __aenter__(self):
        await self.start()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.stop()
